{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FunctionAgent / AgentWorkflow Basic Introduction\n",
    "\n",
    "The `AgentWorkflow` is an orchestrator for running a system of one or more agents. In this example, we'll create a simple workflow with a single `FunctionAgent`, and use that to cover the basic functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "In this example, we will use `OpenAI` as our LLM. For all LLMs, check out the [examples documentation](https://docs.llamaindex.ai/en/stable/examples/llm/openai/) or [LlamaHub](https://llamahub.ai/?tab=llms) for a list of all supported LLMs and how to install/use them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4o-mini\", api_key=\"sk-...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make our agent more useful, we can give it tools/actions to use. In this case, we'll use Tavily to implement a tool that can search the web for information. You can get a free API key from [Tavily](https://tavily.com/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install tavily-python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When creating a tool, its very important to:\n",
    "- give the tool a proper name and docstring/description. The LLM uses this to understand what the tool does.\n",
    "- annotate the types. This helps the LLM understand the expected input and output types.\n",
    "- use async when possible, since this will make the workflow more efficient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tavily import AsyncTavilyClient\n",
    "\n",
    "\n",
    "async def search_web(query: str) -> str:\n",
    "    \"\"\"Useful for using the web to answer questions.\"\"\"\n",
    "    client = AsyncTavilyClient(api_key=\"tvly-...\")\n",
    "    return str(await client.search(query))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the tool and and LLM defined, we can create an `AgentWorkflow` that uses the tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    tools=[search_web],\n",
    "    llm=llm,\n",
    "    system_prompt=\"You are a helpful assistant that can search the web for information.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the Agent\n",
    "\n",
    "Now that our agent is created, we can run it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The current weather in San Francisco is as follows:\n",
      "\n",
      "- **Temperature**: 16.1°C (61°F)\n",
      "- **Condition**: Partly cloudy\n",
      "- **Wind**: 13.6 mph (22.0 kph) from the west\n",
      "- **Humidity**: 64%\n",
      "- **Visibility**: 16 km (9 miles)\n",
      "- **Pressure**: 1017 mb (30.04 in)\n",
      "\n",
      "For more details, you can check the full report [here](https://www.weatherapi.com/).\n"
     ]
    }
   ],
   "source": [
    "response = await agent.run(user_msg=\"What is the weather in San Francisco?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above is the equivalent of the following of using `AgentWorkflow` with a single `FunctionAgent`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import AgentWorkflow\n",
    "\n",
    "workflow = AgentWorkflow(agents=[agent])\n",
    "\n",
    "response = await workflow.run(user_msg=\"What is the weather in San Francisco?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you were creating a workflow with multiple agents, you can pass in a list of agents to the `AgentWorkflow` constructor. Learn more in our [multi-agent workflow example](https://docs.llamaindex.ai/en/stable/understanding/agent/multi_agent/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maintaining State\n",
    "\n",
    "By default, the `FunctionAgent` will maintain stateless between runs. This means that the agent will not have any memory of previous runs.\n",
    "\n",
    "To maintain state, we need to keep track of the previous state. Since the `FunctionAgent` is running in a  `Workflow`, the state is stored in the `Context`. This can be passed between runs to maintain state and history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import Context\n",
    "\n",
    "ctx = Context(agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Nice to meet you, Logan! How can I assist you today?\n"
     ]
    }
   ],
   "source": [
    "response = await agent.run(\n",
    "    user_msg=\"My name is Logan, nice to meet you!\", ctx=ctx\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your name is Logan.\n"
     ]
    }
   ],
   "source": [
    "response = await agent.run(user_msg=\"What is my name?\", ctx=ctx)\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The context is serializable, so it can be saved to a database, file, etc. and loaded back in later. \n",
    "\n",
    "The `JsonSerializer` is a simple serializer that uses `json.dumps` and `json.loads` to serialize and deserialize the context.\n",
    "\n",
    "The `JsonPickleSerializer` is a serializer that uses `pickle` to serialize and deserialize the context. If you have objects in your context that are not serializable, you can use this serializer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import JsonPickleSerializer, JsonSerializer\n",
    "\n",
    "ctx_dict = ctx.to_dict(serializer=JsonSerializer())\n",
    "\n",
    "restored_ctx = Context.from_dict(agent, ctx_dict, serializer=JsonSerializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Yes, I remember your name is Logan.\n"
     ]
    }
   ],
   "source": [
    "response = await agent.run(\n",
    "    user_msg=\"Do you still remember my name?\", ctx=restored_ctx\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming\n",
    "\n",
    "The `AgentWorkflow`/`FunctionAgent` also supports streaming. Since the `AgentWorkflow` is a `Workflow`, it can be streamed like any other `Workflow`. This works by using the handler that is returned from the workflow. There are a few key events that are streamed, feel free to explore below.\n",
    "\n",
    "If you only want to stream the LLM output, you can use the `AgentStream` events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The current weather in Saskatoon is as follows:\n",
      "\n",
      "- **Temperature**: 22.2°C (72°F)\n",
      "- **Condition**: Overcast\n",
      "- **Humidity**: 25%\n",
      "- **Wind Speed**: 6.0 mph (9.7 kph) from the northwest\n",
      "- **Visibility**: 4.8 km\n",
      "- **Pressure**: 1018 mb\n",
      "\n",
      "For more details, you can check the full report [here](https://www.weatherapi.com/)."
     ]
    }
   ],
   "source": [
    "from llama_index.core.agent.workflow import (\n",
    "    AgentInput,\n",
    "    AgentOutput,\n",
    "    ToolCall,\n",
    "    ToolCallResult,\n",
    "    AgentStream,\n",
    ")\n",
    "\n",
    "handler = agent.run(user_msg=\"What is the weather in Saskatoon?\")\n",
    "\n",
    "async for event in handler.stream_events():\n",
    "    if isinstance(event, AgentStream):\n",
    "        print(event.delta, end=\"\", flush=True)\n",
    "        # print(event.response)  # the current full response\n",
    "        # print(event.raw)  # the raw llm api response\n",
    "        # print(event.current_agent_name)  # the current agent name\n",
    "    # elif isinstance(event, AgentInput):\n",
    "    #    print(event.input)  # the current input messages\n",
    "    #    print(event.current_agent_name)  # the current agent name\n",
    "    # elif isinstance(event, AgentOutput):\n",
    "    #    print(event.response)  # the current full response\n",
    "    #    print(event.tool_calls)  # the selected tool calls, if any\n",
    "    #    print(event.raw)  # the raw llm api response\n",
    "    # elif isinstance(event, ToolCallResult):\n",
    "    #    print(event.tool_name)  # the tool name\n",
    "    #    print(event.tool_kwargs)  # the tool kwargs\n",
    "    #    print(event.tool_output)  # the tool output\n",
    "    # elif isinstance(event, ToolCall):\n",
    "    #     print(event.tool_name)  # the tool name\n",
    "    #     print(event.tool_kwargs)  # the tool kwargs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tools and State\n",
    "\n",
    "Tools can also be defined that have access to the workflow context. This means you can set and retrieve variables from the context and use them in the tool or between tools.\n",
    "\n",
    "**Note:** The `Context` parameter should be the first parameter of the tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your name has been set to Logan.\n",
      "Logan\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.workflow import Context\n",
    "\n",
    "\n",
    "async def set_name(ctx: Context, name: str) -> str:\n",
    "    async with ctx.store.edit_state() as ctx_state:\n",
    "        ctx_state[\"state\"][\"name\"] = name\n",
    "    return f\"Name set to {name}\"\n",
    "\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    tools=[set_name],\n",
    "    llm=llm,\n",
    "    system_prompt=\"You are a helpful assistant that can set a name.\",\n",
    "    initial_state={\"name\": \"unset\"},\n",
    ")\n",
    "\n",
    "ctx = Context(agent)\n",
    "\n",
    "response = await agent.run(user_msg=\"My name is Logan\", ctx=ctx)\n",
    "print(str(response))\n",
    "\n",
    "state = await ctx.store.get(\"state\")\n",
    "print(state[\"name\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Human in the Loop\n",
    "\n",
    "Tools can also be defined that involve a human in the loop. This is useful for tasks that require human input, such as confirming a tool call or providing feedback.\n",
    "\n",
    "Using workflow events, we can emit events that require a response from the user. Here, we use the built-in `InputRequiredEvent` and `HumanResponseEvent` to handle the human in the loop, but you can also define your own events.\n",
    "\n",
    "`wait_for_event` will emit the `waiter_event` and wait until it sees the `HumanResponseEvent` with the specified `requirements`. The `waiter_id` is used to ensure that we only send one `waiter_event` for each `waiter_id`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import (\n",
    "    Context,\n",
    "    InputRequiredEvent,\n",
    "    HumanResponseEvent,\n",
    ")\n",
    "\n",
    "\n",
    "async def dangerous_task(ctx: Context) -> str:\n",
    "    \"\"\"A dangerous task that requires human confirmation.\"\"\"\n",
    "\n",
    "    question = \"Are you sure you want to proceed?\"\n",
    "    response = await ctx.wait_for_event(\n",
    "        HumanResponseEvent,\n",
    "        waiter_id=question,\n",
    "        waiter_event=InputRequiredEvent(\n",
    "            prefix=question,\n",
    "            user_name=\"Logan\",\n",
    "        ),\n",
    "        requirements={\"user_name\": \"Logan\"},\n",
    "    )\n",
    "    if response.response == \"yes\":\n",
    "        return \"Dangerous task completed successfully.\"\n",
    "    else:\n",
    "        return \"Dangerous task aborted.\"\n",
    "\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    tools=[dangerous_task],\n",
    "    llm=llm,\n",
    "    system_prompt=\"You are a helpful assistant that can perform dangerous tasks.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The dangerous task has been completed successfully. If you need anything else, feel free to ask!\n"
     ]
    }
   ],
   "source": [
    "handler = agent.run(user_msg=\"I want to proceed with the dangerous task.\")\n",
    "\n",
    "async for event in handler.stream_events():\n",
    "    if isinstance(event, InputRequiredEvent):\n",
    "        response = input(event.prefix).strip().lower()\n",
    "        handler.ctx.send_event(\n",
    "            HumanResponseEvent(\n",
    "                response=response,\n",
    "                user_name=event.user_name,\n",
    "            )\n",
    "        )\n",
    "\n",
    "response = await handler\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In production scenarios, you might handle human-in-the-loop over a websocket or multiple API requests.\n",
    "\n",
    "As mentioned before, the `Context` object is serializable, and this means we can also save the workflow mid-run and restore it later. \n",
    "\n",
    "**NOTE:** Any functions/steps that were in-progress will start from the beginning when the workflow is restored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The dangerous task has been completed successfully. If you need anything else, feel free to ask!\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.workflow import JsonSerializer\n",
    "\n",
    "handler = agent.run(user_msg=\"I want to proceed with the dangerous task.\")\n",
    "\n",
    "input_ev = None\n",
    "async for event in handler.stream_events():\n",
    "    if isinstance(event, InputRequiredEvent):\n",
    "        input_ev = event\n",
    "        break\n",
    "\n",
    "# save the context somewhere for later\n",
    "ctx_dict = handler.ctx.to_dict(serializer=JsonSerializer())\n",
    "\n",
    "# get the response from the user\n",
    "response_str = input(input_ev.prefix).strip().lower()\n",
    "\n",
    "# restore the workflow\n",
    "restored_ctx = Context.from_dict(agent, ctx_dict, serializer=JsonSerializer())\n",
    "\n",
    "handler = agent.run(ctx=restored_ctx)\n",
    "handler.ctx.send_event(\n",
    "    HumanResponseEvent(\n",
    "        response=response_str,\n",
    "        user_name=input_ev.user_name,\n",
    "    )\n",
    ")\n",
    "response = await handler\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
